import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import shap
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import cross_val_score

# Ensure the picture folder exists
def ensure_picture_folder_exists():
    """Ensure the picture folder exists."""
    picture_folder = "./picture"
    if not os.path.exists(picture_folder):
        os.makedirs(picture_folder)
    return picture_folder

# Load data from CSV files
def load_data(medal_data_file, program_data_file):
    """Load data from CSV files."""
    try:
        medal_data = pd.read_csv(medal_data_file, encoding='ISO-8859-1')
        program_data = pd.read_csv(program_data_file, encoding='ISO-8859-1')
        print(f"Data files '{medal_data_file}' and '{program_data_file}' read successfully!")
        return medal_data, program_data
    except FileNotFoundError as e:
        print(f"Error: {e}")
        exit()
    except UnicodeDecodeError as e:
        print(f"Error: {e}")
        exit()

# Clean column names by removing spaces
def clean_column_names(medal_data, program_data):
    """Clean column names by removing spaces."""
    medal_data.columns = medal_data.columns.str.strip()
    program_data.columns = program_data.columns.str.strip()
    return medal_data, program_data

# Check column names to ensure 'Year' column exists
def check_column_names(medal_data, program_data):
    """Check column names to ensure 'Year' column exists."""
    print("medal_data columns:", medal_data.columns)
    print("program_data columns:", program_data.columns)

# Transform 'Year' columns in 'program_data' to rows
def transform_program_data(program_data):
    """Transform 'Year' columns in 'program_data' to rows."""
    program_data_melted = program_data.melt(id_vars=['Sport', 'Discipline', 'Code', 'Sports Governing Body'],
                                            var_name='Year', value_name='Programs')
    return program_data_melted

# Fill missing values and convert categorical features to numeric
def preprocess_data(medal_data, program_data_melted):
    """Fill missing values and convert categorical features to numeric."""
    medal_data.fillna(0, inplace=True)
    program_data_melted.fillna(0, inplace=True)
    medal_data['NOC'] = medal_data['NOC'].astype('category').cat.codes
    program_data_melted['Code'] = program_data_melted['Code'].astype('category').cat.codes
    return medal_data, program_data_melted

# Clean 'Year' column and ensure type consistency
def clean_year_column(program_data_melted, medal_data):
    """Clean 'Year' column and ensure type consistency."""
    program_data_melted['Year'] = program_data_melted['Year'].astype(str).str.replace(r'[^0-9]', '', regex=True).astype(int)
    medal_data['Year'] = medal_data['Year'].astype(int)
    return program_data_melted, medal_data

# Merge medal data with program data
def integrate_data(medal_data, program_data_melted):
    """Merge medal data with program data."""
    merged_data = pd.merge(medal_data, program_data_melted, left_on=['Year', 'NOC'], right_on=['Year', 'Code'], how='left')
    return merged_data

# Extract features and target variables from the merged data
def extract_features(merged_data):
    """Extract features and target variables from the merged data."""
    features = ['Year', 'NOC', 'Gold', 'Silver', 'Bronze', 'Total', 'Programs']
    X = merged_data[features]
    y_gold = merged_data['Gold']  # Target: Predict gold medals
    y_total = merged_data['Total']  # Target: Predict total medals
    return X, y_gold, y_total

# Check data types to ensure all features are numeric
def check_data_types(X):
    """Check data types to ensure all features are numeric."""
    print("Data type check:")
    print(X.dtypes)

# Convert all feature columns to numeric type
def convert_to_numeric(X):
    """Convert all feature columns to numeric type."""
    X = X.apply(pd.to_numeric, errors='coerce')
    return X

# Check for missing values in the features and target variables
def check_missing_values(X, y_gold, y_total):
    """Check for missing values in the features and target variables."""
    print("Missing values check:")
    print(X.isnull().sum())
    print(y_gold.isnull().sum())
    print(y_total.isnull().sum())

# Fill missing values with 0
def handle_missing_values(X, y_gold, y_total):
    """Fill missing values with 0."""
    X.fillna(0, inplace=True)
    y_gold.fillna(0, inplace=True)
    y_total.fillna(0, inplace=True)
    return X, y_gold, y_total

# Train Random Forest Regressor models and evaluate their performance
def train_models(X, y_gold, y_total):
    """Train Random Forest Regressor models and evaluate their performance."""
    model_gold = RandomForestRegressor(n_estimators=100, random_state=42)
    model_total = RandomForestRegressor(n_estimators=100, random_state=42)

    # Cross-validation to evaluate model performance
    cv_scores_gold = cross_val_score(model_gold, X, y_gold, cv=5, scoring='neg_mean_squared_error', error_score='raise')
    cv_scores_total = cross_val_score(model_total, X, y_total, cv=5, scoring='neg_mean_squared_error', error_score='raise')

    print(f"Cross-validation MSE for gold medal prediction: {-cv_scores_gold.mean():.2f}")
    print(f"Cross-validation MSE for total medal prediction: {-cv_scores_total.mean():.2f}")

    # Train models
    model_gold.fit(X, y_gold)
    model_total.fit(X, y_total)

    return model_gold, model_total

# Visualize the model performance and feature importance
def visualize_model_performance(model_gold, X, y_gold, y_total, picture_folder):
    """Visualize the model performance and feature importance."""
    features = X.columns
    # 15.1 Plot feature importance for gold medal prediction
    plt.figure(figsize=(10, 6))
    plt.barh(features, model_gold.feature_importances_, color='skyblue')
    plt.title('Feature Importance for Gold Medal Prediction')
    plt.xlabel('Importance')
    plt.ylabel('Feature')
    plt.grid(axis='x')  # Add vertical grid lines
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.1.png")
    plt.show()

    # 15.2 Calculate predictions
    y_pred_gold = model_gold.predict(X)  # Predicted gold medals
    y_pred_total = model_total.predict(X)  # Predicted total medals

    # 15.3 Plot actual vs predicted gold medals
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x=y_gold, y=y_pred_gold, color='blue', alpha=0.6)
    plt.plot([min(y_gold), max(y_gold)], [min(y_gold), max(y_gold)], color='red', linestyle='--')
    plt.title('Actual vs Predicted Gold Medals')
    plt.xlabel('Actual Gold Medals')
    plt.ylabel('Predicted Gold Medals')
    plt.grid(True)  # Add grid lines
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.2.png")
    plt.show()

    # 15.4 Plot actual vs predicted total medals
    plt.figure(figsize=(10, 6))
    sns.scatterplot(x=y_total, y=y_pred_total, color='green', alpha=0.6)
    plt.plot([min(y_total), max(y_total)], [min(y_total), max(y_total)], color='red', linestyle='--')
    plt.title('Actual vs Predicted Total Medals')
    plt.xlabel('Actual Total Medals')
    plt.ylabel('Predicted Total Medals')
    plt.grid(True)  # Add grid lines
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.3.png")
    plt.show()

# Perform SHAP analysis for model explainability
def shap_analysis(model_gold, model_total, X, picture_folder):
    """Perform SHAP analysis for model explainability."""
    # Initialize SHAP explainer
    explainer_gold = shap.TreeExplainer(model_gold)
    shap_values_gold = explainer_gold.shap_values(X)

    explainer_total = shap.TreeExplainer(model_total)
    shap_values_total = explainer_total.shap_values(X)

    # 16.1 Plot SHAP values summary for gold medal prediction
    plt.figure(figsize=(10, 6))
    shap.summary_plot(shap_values_gold, X, feature_names=X.columns, show=False)
    plt.title('SHAP Summary Plot for Gold Medal Prediction')
    plt.grid(True)  # Add grid lines
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.4.png")
    plt.show()

    # 16.2 Plot SHAP values summary for total medal prediction
    plt.figure(figsize=(10, 6))
    shap.summary_plot(shap_values_total, X, feature_names=X.columns, show=False)
    plt.title('SHAP Summary Plot for Total Medal Prediction')
    plt.grid(True)  # Add grid lines
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.5.png")
    plt.show()

    # 16.3 SHAP force plot for a single sample (gold medal prediction)
    sample_idx = 0  # Select the first sample
    shap.force_plot(explainer_gold.expected_value, shap_values_gold[sample_idx, :], X.iloc[sample_idx, :], feature_names=X.columns)

# Analyze the impact of the number and type of programs on medal distribution
def analyze_medal_distribution(merged_data, picture_folder):
    """Analyze the impact of the number and type of programs on medal distribution."""
    # 17.1 Aggregate medal counts by sport and country
    sport_medals = merged_data.groupby(['NOC', 'Sport']).agg({'Gold': 'sum', 'Silver': 'sum', 'Bronze': 'sum', 'Total': 'sum'}).reset_index()

    # 17.2 Plot medal distribution by sport and country
    plt.figure(figsize=(18, 10))
    sns.barplot(x='NOC', y='Total', hue='Sport', data=sport_medals, palette='viridis', width=0.9)

    # Limit legend to show only the top 10 sports
    handles, labels = plt.gca().get_legend_handles_labels()
    plt.legend(handles=handles[:10], labels=labels[:10], title="Sport", bbox_to_anchor=(1.05, 1), loc='upper left')

    plt.title('Medal Distribution by Sport and Country')
    plt.xlabel('Country Code')
    plt.ylabel('Total Medals')
    plt.xticks(rotation=90)
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.6.png")
    plt.show()

# Analyze the importance of programs to countries' total medals
def analyze_program_importance(merged_data, picture_folder):
    """Analyze the importance of programs to countries' total medals."""
    # 18.1 Calculate each sport's contribution to a country's total medals
    program_importance = merged_data.groupby(['NOC', 'Sport', 'Discipline']).agg({'Total': 'sum'}).reset_index()
    program_importance['Contribution'] = program_importance.groupby('NOC')['Total'].transform(lambda x: x / x.sum())

    # 18.2 Plot sport contribution to total medals by country
    plt.figure(figsize=(18, 10))
    sns.barplot(x='NOC', y='Contribution', hue='Sport', data=program_importance, palette='viridis', width=0.9)

    # Limit legend to show only the top 10 sports
    handles, labels = plt.gca().get_legend_handles_labels()
    plt.legend(handles=handles[:10], labels=labels[:10], title="Sport", bbox_to_anchor=(1.05, 1), loc='upper left')

    plt.title('Contribution of Sports to Total Medals by Country')
    plt.xlabel('Country Code')
    plt.ylabel('Contribution to Total Medals')
    plt.xticks(rotation=90)
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.7.png")
    plt.show()

# Analyze the impact of being a host country on the number of programs and types of sports
def analyze_host_impact(merged_data, picture_folder):
    """Analyze the impact of being a host country on the number of programs and types of sports."""
    # 19.1 Assume the host country for the 2028 Olympics is a specific country (e.g., country code 0)
    host_country = 0
    host_data = merged_data[merged_data['NOC'] == host_country]
    non_host_data = merged_data[merged_data['NOC'] != host_country]

    # 19.2 Compare the number of programs between host and non-host countries
    plt.figure(figsize=(12, 8))
    sns.boxplot(x='NOC', y='Programs', data=merged_data, palette='Set2')
    plt.title('Number of Programs by Country (Host vs Non-Host)')
    plt.xlabel('Country Code')
    plt.ylabel('Number of Programs')
    plt.xticks(rotation=90)
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.8.png")
    plt.show()

    # 19.3 Compare the types of sports between host and non-host countries
    host_sports = host_data['Sport'].value_counts().reset_index()
    host_sports.columns = ['Sport', 'Count']

    non_host_sports = non_host_data['Sport'].value_counts().reset_index()
    non_host_sports.columns = ['Sport', 'Count']

    # Plot comparison of sports between host and non-host countries
    plt.figure(figsize=(14, 8))
    sns.barplot(x='Sport', y='Count', data=host_sports, color='blue', label='Host Country')
    sns.barplot(x='Sport', y='Count', data=non_host_sports, color='green', label='Non-Host Countries', alpha=0.6)

    plt.title('Comparison of Sports between Host and Non-Host Countries')
    plt.xlabel('Sport')
    plt.ylabel('Count')

    # Set y-axis interval to 5
    plt.yticks(np.arange(0, max(host_sports['Count'].max(), non_host_sports['Count'].max()) + 1, step=5))

    plt.xticks(rotation=90)
    plt.legend()
    plt.tight_layout()
    plt.savefig(f"{picture_folder}/1.4.9.png")
    plt.show()

# Main execution
if __name__ == "__main__":
    # Ensure the picture folder exists
    picture_folder = ensure_picture_folder_exists()

    # Step 1: Load data
    medal_data_file = r'summerOly_medal_counts.csv'
    program_data_file = r'summerOly_programs.csv'
    medal_data, program_data = load_data(medal_data_file, program_data_file)

    # Step 2: Clean column names
    medal_data, program_data = clean_column_names(medal_data, program_data)

    # Step 3: Check column names
    check_column_names(medal_data, program_data)

    # Step 4: Transform program data
    program_data_melted = transform_program_data(program_data)

    # Step 5: Data preprocessing
    medal_data, program_data_melted = preprocess_data(medal_data, program_data_melted)

    # Step 6: Clean year column
    program_data_melted, medal_data = clean_year_column(program_data_melted, medal_data)

    # Step 7: Data integration
    merged_data = integrate_data(medal_data, program_data_melted)

    # Step 8: Feature extraction
    X, y_gold, y_total = extract_features(merged_data)

    # Step 9: Check data types
    check_data_types(X)

    # Step 10: Convert to numeric
    X = convert_to_numeric(X)

    # Step 11: Check for missing values
    check_missing_values(X, y_gold, y_total)

    # Step 12: Handle missing values
    X, y_gold, y_total = handle_missing_values(X, y_gold, y_total)

    # Step 13: Model training
    model_gold, model_total = train_models(X, y_gold, y_total)

    # Step 14: Model visualization
    visualize_model_performance(model_gold, X, y_gold, y_total, picture_folder)

    # Step 15: SHAP analysis
    shap_analysis(model_gold, model_total, X, picture_folder)

    # Step 16: Analyze medal distribution
    analyze_medal_distribution(merged_data, picture_folder)

    # Step 17: Analyze program importance
    analyze_program_importance(merged_data, picture_folder)

    # Step 18: Analyze host impact
    analyze_host_impact(merged_data, picture_folder)